
import os
import pickle
import random

import numpy as np
import copy

import torch
import time

from .model import HierarchicalDQN
from .memory_buffer import ReplayMemory
from .utils import state_to_features

ACTIONS = ['UP', 'RIGHT', 'DOWN', 'LEFT', 'WAIT', 'BOMB']

def setup(self):
    """
    Setup your code. This is called once when loading each agent.
    Make sure that you prepare everything such that act(...) can be called.

    When in training mode, the separate `setup_training` in train.py is called
    after this method. This separation allows you to share your trained agent
    with other students, without revealing your training code.

    In this example, our model is a set of probabilities over actions
    that are is independent of the game state.

    :param self: This object is passed to all callbacks and you can set arbitrary values.
    """
    if self.train or not os.path.isfile("my-saved-model.pt"):
        self.logger.info("Setting up model from scratch.")
        self.model = HierarchicalDQN()
        self.memory = ReplayMemory(1024)
        self.model.switch_train_mode()
    else:
        self.logger.info("Loading model from saved state.")
        self.model = HierarchicalDQN()
        self.memory = ReplayMemory(1024)
        self.model.load_model('my-saved-model.pt')
        self.model.switch_eval_mode()

def act(self, game_state: dict) -> str:
    """
    Your agent should parse the input, think, and take a decision.
    When not in training mode, the maximum execution time for this method is 0.5s.

    :param self: The same object that is passed to all of your callbacks.
    :param game_state: The dictionary that describes everything on the board.
    :return: The action to take as a string.
    """
    
    if not self.train:
        self.memory.push(game_state)

    observation = state_to_features(game_state)
    state = {
        'observation':observation,
        'game_state':game_state
    }
    action = self.memory.action(state,self.model)

    # self.logger.debug("Querying model for action.")
    return action

